#!/usr/bin/env python

import sys
import os
import argparse
from itertools import chain

parser = argparse.ArgumentParser(description= """
DESCRIPTION
    Parses a list of file names to return those matching a range of library IDs.

MEMO:
    If using pyhton range() func remember that the end of the range is not returned!
    E.g. range(1, 4) >>> [1, 2, 3]

EXAMPLE
    find fastq -name 'mjb0[0-9][0-9]' | idFind.py --initials mjb --range '[range(1, 10), 12, 14]'

""", formatter_class= argparse.RawTextHelpFormatter)

parser.add_argument('--infiles', '-i',
                   required= False,
                   nargs= '+',
                   help='''List of files to parse. Default to stdin.
                   ''')

parser.add_argument('--initials', '-I',
                   required= True,
                   help='''Initials to get from beginning of filename.
                   ''')

parser.add_argument('--range', '-r',
                   required= False,
                   help='''A string evaluated by python as list or tuple containing
the numeric IDs to return. 
                   ''')

def flatten(*args):
    for x in args:
        if hasattr(x, '__iter__'):
            for y in flatten(*x):
                yield y
        else:
            yield x

def evalRange(id_range):
    """Evaluates the string id_range to return the integers
    corresponding to library IDs.
    - The elements of the evaluated string must be integers.
    - Duplicates are removed
    """
    ids= eval(id_range)
    ids= list(flatten(ids))
    for x in ids:
        try:
            int(x)
        except ValueError:
            return(None)
    idsNoDup= []
    for x in ids:
        if x not in idsNoDup:
            idsNoDup.append(x)
    return(idsNoDup)

def getFilesForInitials(fileList, initials):
    """Given a list of files (fileList), return the files that begin with
    initials. Files are returned witht the same path.
    """
    initialFiles= []
    for f in fileList:
        fpath, fname= os.path.split(f)
        if fname.startswith(initials):
            initialFiles.append(os.path.join(fpath, fname))
    return(initialFiles)

def getNumericID(fileList, initials):
    """Given the list of files, return the numeric identifier
    of each ID as tuple (file, id)
    """
    numericID= []
    for x in fileList:
        fpath, fname= os.path.split(x)
        if not fname.startswith(initials):
            sys.exit('Filename does not start with requested initials %s, %s' %(x, initials))
        libid= fname[len(initials):]
        try:
            intID= int(libid.split('_')[0])
        except ValueError:
            continue
            # sys.exit('Could not extract valid numeric identifier for %s' %(x))
        numericID.append((x, intID))
    return((numericID))

def getFilesForInt(fileTuples, fileInts):
    """Given the list of tuples [(filename, library-int)]
    return the files whose library-int is in the list fileInts.
    All the files are returned if fileInts is [].
    """
    selectedFiles= []
    for f, i in fileTuples:
        if i in fileInts or fileInts == []:
            selectedFiles.append(f)
    return(selectedFiles)
    
def main():
    args = parser.parse_args()
    if not args.range:
        int_range= []
    else:
        int_range= evalRange(args.range)
    if not args.infiles or args.infiles == ['-']:
        infiles= sys.stdin.readlines()
        infiles= [x.strip() for x in infiles]
    else:
        infiles= args.infiles
    initFiles= getFilesForInitials(infiles, args.initials)
    initFiles= getNumericID(initFiles, args.initials)
    selectFiles= getFilesForInt(initFiles, int_range)
    for f in selectFiles:
        print(f)

if __name__ == '__main__':
    main()
    sys.exit()